#!/usr/bin/env python

# 屏蔽　ssh 攻击 v1.0

import logging
import logging.config
import config_default
import data_api
from datetime import *
import re
import subprocess
from io import StringIO

logging.config.dictConfig(config_default.configs['log'])

'''
持久化格式
{
  record_time: '2014-01-16 16:46:00',
  fail_dict: {
    '94.156.65.21': {
      time: '2024-01-17 16:23:29',
      count: 6
    },
    '94.156.65.21': {
      time: '2024-01-17 16:23:29',
      count: 6
    }
  },
  lock_ip : [
    { ip: '172.232.159.144', time: '2024-01-17 09:41:50' },
    { ip: '172.232.159.144', time: '2024-01-17 09:41:50' }
  ]
}
'''

def update_record_time(time):
  data=data_api.read_data()
  data['record_time'] = time.isoformat(' ', 'seconds')
  data_api.write_data(data)

# 取上一次登记记录时间，如果没有记录，取当前时间向前6小时
def last_record_time():
  data = data_api.read_data()
  exsit= 'record_time' in data
  return datetime.fromisoformat(data['record_time']) if exsit else datetime.now() - timedelta(hours=6)

# 解析由`lastb -Fia -n N` 返回的单行数据，如有效返回ip,否则返回None
def parse_lastb(line, last_record_time):
  list = line.strip().split()
  ip = list[-1]
  #print('ip -> %s, time -> %s' % (list[-1], ' '.join(list[2:7])))
  if not re.fullmatch(r'\d+(\.\d+){3}', ip):
    logging.error('line %s -- ip not match' % line)
    return None
  start = 2
  if 'ssh:notty' in list:
    start = list.index('ssh:notty') + 1
  time = datetime.strptime(' '.join(list[start:start + 5]), '%a %b %d %H:%M:%S %Y')
  if time > last_record_time:
    return ip
  else:
    return None

def update_fail_dict(lastb_list_str, now):
  stream = StringIO(lastb_list_str)
  last_time = last_record_time()
  now_str = now.isoformat(' ', 'seconds')
  data = data_api.read_data()
  fail_dict = data.get('fail_dict', {})
  while True:
    line = stream.readline().strip()
    if len(line) <= 0:
      break
    ip = parse_lastb(line, last_time)
    if not ip:
      continue
    if ip in fail_dict:
      info = fail_dict[ip]
      if (now - datetime.fromisoformat(info['time'])).total_seconds() > 3600:
        info['count'] = 1
        info['time'] = now_str
      else:
        info['count'] = info['count'] + 1
    else:
      fail_dict[ip] = { 'time': now_str, 'count': 1 }
  #print(fail_dict)
  data['fail_dict'] = fail_dict
  data_api.write_data(data)
  update_record_time(now)
 

# 防火墙屏蔽　ip
def firewall_lock_ip(now, fail_times):
  data = data_api.read_data()
  fail_dict = data.get('fail_dict', {})
  remove_keys = []
  lock_ips = []
  for (ip, info) in fail_dict.items():
    if info['count'] > fail_times:
      lock_ips.append(ip)
      remove_keys.append(ip)
    else:
      duration = (now - datetime.fromisoformat(info['time'])).total_seconds()
      if duration > 24 * 3600:
        remove_keys.append(ip)
  # lock ip
  now_str = now.isoformat(' ', 'seconds')
  lock_ip_list = data.get('lock_ip', [])
  for ip in lock_ips:
    logging.info('lock ip --> %s ' % ip)
    centos7_firewall_drop_ip(ip)
    lock_ip_list.append({ 'time': now_str, 'ip': ip })
  # remove
  for ip in remove_keys:
    del fail_dict[ip]
  data['lock_ip'] = lock_ip_list
  data['fail_dict'] = fail_dict
  data_api.write_data(data)

# 超过锁定时长的 ip 解除
def firewall_unlock(now, lock_seconds):
  data = data_api.read_data()
  lock_ip_list = data.get('lock_ip', [])
  list = []
  for info in lock_ip_list:
    duration = (now - datetime.fromisoformat(info['time'])).total_seconds()
    if duration > lock_seconds:
      logging.info('unlock ip --> %s, lock time - %s' % (info['ip'], info['time']))
      centos7_firewall_undo_drop_ip(info['ip'])
    else:
      list.append(info)
  data['lock_ip'] = list
  data_api.write_data(data)

def centos7_firewall_drop_ip(ip):
  cmd = 'firewall-cmd --add-rich-rule="rule family=ipv4 source address=%s drop"'
  subprocess.run(cmd % ip, shell=True, stdout=subprocess.DEVNULL)

def centos7_firewall_undo_drop_ip(ip):
  cmd = 'firewall-cmd --remove-rich-rule="rule family=ipv4 source address=%s drop"'
  subprocess.run(cmd % ip, shell=True, stdout=subprocess.DEVNULL)


# 主函数启动入口
def main():
  try:
    # 如果登录失败，会被锁定的时长　lock_seconds 秒
    lock_seconds = 1 * 24 * 3600
    # 允许失败次数
    fail_times = 3
    now = datetime.now()
    lastb_record = subprocess.run(['lastb', '-Fia', '-30'], capture_output=True).stdout.decode('utf-8')
    update_fail_dict(lastb_record, now)
    firewall_lock_ip(now, fail_times)
    firewall_unlock(now, lock_seconds)
  except BaseException as e:
    logging.error(e)


def test_unlock():
  firewall_unlock(datetime.now(), 10*60)


def test_lock_ip():
  firewall_lock_ip(datetime.now(), 3)

def test_lastb_record():
  out = subprocess.run(['ssh', 'root@192.168.50.30', 'lastb -Fia -30'], capture_output=True).stdout.decode('utf-8')
  update_fail_dict(out, datetime.now())


def test_parse_lastb():
  line = 'user     ssh:notty    Wed Jan 21 09:41:50 2024 - Wed Jan 17 09:41:50 2024  (00:00)     94.156.65.21'
  #line = '         ssh:notty    Sun Jan 28 06:14:57 2024 - Sun Jan 28 06:14:57 2024  (00:00)     123.160.221.137'
  #line = 'btmp begins Mon Jan  1 03:34:05 2024'
  last_time = last_record_time()
  print(parse_lastb(line, last_time))

def test_update_recordtime():
  now = datetime.now()
  update_record_time(now)


# 时间测试
def test_datetime():
  now = datetime.now()
  logging.debug(now)
  print(now.isoformat(sep=' ', timespec='seconds'))
  print(now.isoformat(' ', 'seconds'))
  record = datetime.fromisoformat('2024-01-16 17:08:39')
  print(record)
  print(now - record)
  print((now - record).total_seconds())
  print(datetime.strptime('Wed Jan 17 09:41:50 2024', '%a %b %d %H:%M:%S %Y'))

if __name__ == '__main__':
  main()
  #test_datetime()
  #test_update_recordtime()
  #print(last_record_time())
  #test_parse_lastb()
  #test_lastb_record()
  #test_unlock()
